#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <random>
#include <systemc>
#include <tlm>
#include <tlm_utils/simple_initiator_socket.h>
#include <tlm_utils/simple_target_socket.h>

using namespace sc_core;
using namespace sc_dt;
using namespace tlm_utils;
using namespace std;

const sc_time DELAY = sc_time(2, SC_NS);

/**

                                                  +---------+
                                         .--->[ts]|  Slave  |
                                        |         +---------+
                                        |
                                        |         +---------+
+--------+            +--------+        |---->[ts]|  Slave  |
| Master |[is]--->[ts]| Router |[is]--->|         +---------+
+--------+            +--------+        |
                                        |         +---------+
                                        |---->[ts]|  Slave  |
                                        |         +---------+
                                        |
                                        |         +---------+
                                         '--->[ts]|  Slave  |
                                                  +---------+
*/

template<size_t N_TARGETS>
SC_MODULE(Master) {
    // Random number generator
    std::mt19937 gen;
    std::uniform_int_distribution<uint8_t> offset_dist;
    std::uniform_int_distribution<uint8_t> index_dist; // 0 - 63
    std::uniform_int_distribution<uint8_t> rwop_dist; // 0: read, 1: write
    // Initiator socket
    simple_initiator_socket<Master> socket;
    // DMI
    bool dmi_ptr_valid = false;
    tlm::tlm_dmi dmi_data;

    SC_CTOR(Master) :
          socket("initiator_socket"),
          gen(0),
          offset_dist(0, 64 - 1),
          index_dist(0, N_TARGETS - 1),
          rwop_dist(0, 1),
          dmi_ptr_valid(false) {
        // Register callbacks
        SC_THREAD(process);
        socket.register_invalidate_direct_mem_ptr(this, &Master::invalidate_direct_mem_ptr);
    }

    void process() {
        sc_time delay = DELAY;
        tlm::tlm_generic_payload payload;
        while (true) {
            uint32_t addr = index_dist(gen) << 8 | offset_dist(gen) << 2;
            uint8_t rwop = rwop_dist(gen);
            uint32_t data = 0;

            // Prepare transaction
            tlm::tlm_command cmd = (rwop == 0) ? tlm::TLM_READ_COMMAND : tlm::TLM_WRITE_COMMAND;
            if (cmd == tlm::TLM_WRITE_COMMAND) { data = 0xFF000000 | addr; }

            // Use DMI if available
            sc_dt::uint64 req_addr = sc_dt::uint64 {addr};
            printf("%s:\t%s REQ addr=0x%04X at time %s delay %s\n",
                    this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                    addr, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            bool in_range = req_addr >= dmi_data.get_start_address() && req_addr <= dmi_data.get_end_address();
            if (dmi_ptr_valid && in_range) {
                auto dmi_addr = dmi_data.get_dmi_ptr() + (req_addr - dmi_data.get_start_address());
                // DMI Read
                if (cmd == tlm::TLM_READ_COMMAND) {
                    assert(dmi_data.is_read_allowed());
                    memcpy(&data, dmi_addr, 4);
                    delay = dmi_data.get_read_latency();
                }
                // DMI Write
                else if (cmd == tlm::TLM_WRITE_COMMAND) {
                    assert(dmi_data.is_write_allowed());
                    memcpy(dmi_addr, &data, 4);
                    delay = dmi_data.get_write_latency();
                }

                // Print transaction details
                printf("%s:\t%s RSP addr=0x%04X, data=0x%08X at time %s delay %s [DMI]\n",
                    this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                    addr, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            }
            else {
                // Set up the payload
                payload.set_command(cmd);
                payload.set_address(addr);
                payload.set_data_ptr(reinterpret_cast<unsigned char*>(&data));
                payload.set_data_length(sizeof(data));
                payload.set_byte_enable_ptr(0);
                payload.set_dmi_allowed(false);
                payload.set_response_status(tlm::TLM_INCOMPLETE_RESPONSE);

                // Send the transaction
                socket->b_transport(payload, delay);

                // Check response status
                if (payload.is_response_error()) {
                    SC_REPORT_ERROR("TLM-2", "Response error from b_transport");
                }

                // Check DMI hint
                if (payload.is_dmi_allowed()) {
                    dmi_data.init();
                    // Reuse the same payload to get DMI info.
                    // Reset the address because it may have been modified by the interconnect on the previous transport call.
                    payload.set_address(addr); // Reset address
                    dmi_ptr_valid = socket->get_direct_mem_ptr(payload, dmi_data);
                }

                // Print transaction details
                printf("%s:\t%s RSP addr=0x%04X, data=0x%08X at time %s delay %s\n",
                    this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                    addr, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            }

            // Wait for the delay
            wait(delay);
        }
    }

    // TLM-2.0 backward DMI method
    void invalidate_direct_mem_ptr(sc_dt::uint64 start_range, sc_dt::uint64 end_range) {
        // For simplicity, invalidate the entire DMI region, ignore the range
        printf("%s: DMI pointer invalidated for [0x%llX, 0x%llX] at time %s\n",
               this->name(), start_range, end_range, sc_time_stamp().to_string().c_str());
        dmi_ptr_valid = false;
    }
};

template<size_t N_TARGETS>
SC_MODULE(Router) {
    // Target socket
    tlm_utils::simple_target_socket<Router> target_socket;

    // Initiator sockets
    tlm_utils::simple_initiator_socket_tagged<Router>* initiator_sockets[N_TARGETS];

    SC_CTOR(Router) : target_socket("target_socket") {
        // Register callbacks for incoming interface method calls
        target_socket.register_b_transport(this, &Router::b_transport);
        target_socket.register_get_direct_mem_ptr(this, &Router::get_direct_mem_ptr);

        for (unsigned int i = 0; i < N_TARGETS; i++) {
            initiator_sockets[i] = new tlm_utils::simple_initiator_socket_tagged<Router>((string("initiator_socket_") + to_string(i)).c_str());

            // Register callbacks for incoming interface method calls, including tags
            initiator_sockets[i]->register_invalidate_direct_mem_ptr(
                this, &Router::invalidate_direct_mem_ptr, i);
        }
    }

    // TLM-2.0 blocking transport method
    void b_transport(tlm::tlm_generic_payload &payload, sc_time &delay) {
        sc_dt::uint64 address = payload.get_address();
        sc_dt::uint64 masked_address;
        size_t target_index = decode_address(address, masked_address);
        payload.set_address(masked_address);
        // Forward the payload to the appropriate initiator socket
        (*(initiator_sockets[target_index]))->b_transport(payload, delay);
    }

    // TLM-2.0 forward DMI method
    bool get_direct_mem_ptr(tlm::tlm_generic_payload &payload, tlm::tlm_dmi &dmi_data) {
        // Prepare address
        sc_dt::uint64 masked_address;
        size_t target_index = decode_address(payload.get_address(), masked_address);
        payload.set_address(masked_address);

        // Forward DMI request to appropriate target
        bool status = (*(initiator_sockets[target_index]))->get_direct_mem_ptr(payload, dmi_data);

        // Calculate DMI address of target in system address space
        dmi_data.set_start_address(compose_address(target_index, dmi_data.get_start_address()));
        dmi_data.set_end_address(compose_address(target_index, dmi_data.get_end_address()));

        return status;
    }

    // TLM-2.0 tagged backward DMI method
    // In this example, invalidate_direct_mem_ptr is not used.
    void invalidate_direct_mem_ptr(int id, sc_dt::uint64 start_range, sc_dt::uint64 end_range) {
        sc_dt::uint64 bw_start_range = compose_address(id, start_range);
        sc_dt::uint64 bw_end_range = compose_address(id, end_range);
        // Call target socket's invalidate_direct_mem_ptr
        printf("%s: Invalidate DMI for target %d range [0x%llX, 0x%llX] at time %s\n",
               this->name(), id, bw_start_range, bw_end_range, sc_time_stamp().to_string().c_str());
        target_socket->invalidate_direct_mem_ptr(bw_start_range, bw_end_range);
    }

    // Decode address to determine target and masked address
    size_t decode_address(sc_dt::uint64 address, sc_dt::uint64 &masked_address) {
        masked_address = address & 0xFF;
        return static_cast<size_t>((address >> 8) & 0x3);
    }
    // Compose address from target index and slave address
    sc_dt::uint64 compose_address(size_t target_index, sc_dt::uint64 address) {
        return (target_index << 8) | (address & 0xFF);
    }
};

struct Slave: sc_module {
    // ID
    size_t id;
    // Target socket
    simple_target_socket<Slave> socket;
    // Memory
    uint8_t mem[256];
    // Random number generator
    std::mt19937 gen;
    std::uniform_int_distribution<uint8_t> dist; // 0-255

    Slave(sc_module_name _name, size_t _id)
        : sc_module(_name), id(_id), socket("target_socket"), gen(1), dist(0, 255) {
        // Register callbacks
        socket.register_b_transport(this, &Slave::b_transport);
        socket.register_get_direct_mem_ptr(this, &Slave::get_direct_mem_ptr);
        // Initialize memory
        uint32_t* base = reinterpret_cast<uint32_t*>(&mem[0]);
        for (int i = 0; i < 64; ++i) {
            base[i] = 0xAA000000 | (id << 8) | (i << 2);
        }
    }

    void b_transport(tlm::tlm_generic_payload& payload, sc_time& delay) {
        // Extract command, address, and data pointer
        tlm::tlm_command cmd = payload.get_command();
        uint64_t addr = payload.get_address();
        unsigned char* data_ptr = payload.get_data_ptr();
        unsigned int data_length = payload.get_data_length();

        // Check address and data length
        if (addr >= 256 || data_length != 4) {
            payload.set_response_status(tlm::TLM_ADDRESS_ERROR_RESPONSE);
            return;
        }

        // Operate on memory
        if (cmd == tlm::TLM_READ_COMMAND) {
            memcpy(data_ptr, &mem[addr], data_length);
            printf("%s:\tR     addr=0x%04X, data=0x%08X at time %s\n",
                   this->name(), (unsigned int)addr, *((uint32_t*)data_ptr), sc_time_stamp().to_string().c_str());
        }
        else if (cmd == tlm::TLM_WRITE_COMMAND) {
            memcpy(&mem[addr], data_ptr, data_length);
            printf("%s:\tW     addr=0x%04X, data=0x%08X at time %s\n",
                   this->name(), (unsigned int)addr, *((uint32_t*)data_ptr), sc_time_stamp().to_string().c_str());
        }

        // Set delay
        delay = DELAY * 4;

        // Set response status
        payload.set_response_status(tlm::TLM_OK_RESPONSE);

        // Set DMI hint
        payload.set_dmi_allowed(true);
    }

    // TLM-2.0 forward DMI method
    bool get_direct_mem_ptr(tlm::tlm_generic_payload& payload, tlm::tlm_dmi& dmi_data) {
        // Set DMI properties
        dmi_data.allow_read_write();
        dmi_data.set_dmi_ptr(reinterpret_cast<unsigned char*>(&mem[0]));
        dmi_data.set_start_address(0);
        dmi_data.set_end_address(255);
        dmi_data.set_read_latency(DELAY);
        dmi_data.set_write_latency(DELAY);
        return true;
    }

};

int sc_main(int argc, char **argv)
{
    const size_t N_TARGETS = 4;

    // Instantiate modules
    Master<N_TARGETS>* master;
    Slave*  slaves[N_TARGETS];
    Router<N_TARGETS>* router;

    router = new Router<N_TARGETS> {"router"};
    master = new Master<N_TARGETS> {"master"};
    for (size_t i = 0; i < N_TARGETS; ++i) {
        slaves[i] = new Slave(("slave" + to_string(i)).c_str(), i);
    }

    // Bind sockets
    master->socket.bind(router->target_socket);
    for (size_t i = 0; i < N_TARGETS; ++i) {
        router->initiator_sockets[i]->bind(slaves[i]->socket);
    }

    // Start simulation
    sc_start(1000, SC_NS);
    return 0;
}
